"""
获取并保存上证和深证股票代码的脚本
"""

import time
import h5py
import akshare as ak
import numpy as np
import pandas as pd

from datetime import datetime
from yfinance.exceptions import YFTzMissingError
# from finrl.config_tickers import DOW_30_TICKER
# import os
# os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
def get_shanghai_tickers(max_retries=3):
    """
    获取上海交易所所有股票代码
    
    参数:
    max_retries : int
        最大重试次数，默认为3次
        
    返回:
    list: 上海交易所股票代码列表
    """
    for attempt in range(max_retries):
        try:
            # 获取上海A股实时行情
            sh_stocks = ak.stock_sh_a_spot_em()
            # 转换为FinRL格式
            sh_tickers = [f"{code}.XSHG" for code in sh_stocks['代码']]
            print(f"成功获取上海交易所 {len(sh_tickers)} 只股票代码")
            return sh_tickers
        except Exception as e:
            print(f"第 {attempt + 1} 次尝试获取上海股票代码失败: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # 指数退避
            else:
                print("获取上海股票代码最终失败，返回空列表")
                return []


def get_shenzhen_tickers(max_retries=3):
    """
    获取深圳交易所所有股票代码
    
    参数:
    max_retries : int
        最大重试次数，默认为3次
        
    返回:
    list: 深圳交易所股票代码列表
    """
    for attempt in range(max_retries):
        try:
            # 获取深圳A股实时行情
            sz_stocks = ak.stock_sz_a_spot_em()
            # 转换为FinRL格式
            sz_tickers = [f"{code}.XSHE" for code in sz_stocks['代码']]
            print(f"成功获取深圳交易所 {len(sz_tickers)} 只股票代码")
            return sz_tickers
        except Exception as e:
            print(f"第 {attempt + 1} 次尝试获取深圳股票代码失败: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # 指数退避
            else:
                print("获取深圳股票代码最终失败，返回空列表")
                return []

def saveTickers():
    try:
        # 分别获取上海和深圳的股票代码
        print("正在获取上海交易所股票代码...")
        sh_tickers = get_shanghai_tickers()
        time.sleep(1)  # 每次请求后等待1秒

        print("正在获取深圳交易所股票代码...")
        sz_tickers = get_shenzhen_tickers()

        # 检查是否成功获取到数据
        if not sh_tickers and not sz_tickers:
            print("未能获取到任何股票代码，程序退出")
            return

        # 将上海和深圳的股票代码进行拼接保存到h5文件
        print("正在保存股票代码到h5文件...")
        with h5py.File('stock_ticker.h5', 'w') as h5_file:
            # 保存上海股票代码
            sh_dt = h5py.string_dtype(encoding='utf-8')
            sh_dataset = h5_file.create_dataset(
                'sh_tickers',
                data=np.array(sh_tickers, dtype=sh_dt),
                dtype=sh_dt
            )

            # 保存深圳股票代码
            sz_dt = h5py.string_dtype(encoding='utf-8')
            sz_dataset = h5_file.create_dataset(
                'sz_tickers',
                data=np.array(sz_tickers, dtype=sz_dt),
                dtype=sz_dt
            )

            # 添加文件信息
            h5_file.attrs['sh_count'] = len(sh_tickers)
            h5_file.attrs['sz_count'] = len(sz_tickers)
            h5_file.attrs['total_count'] = len(sh_tickers) + len(sz_tickers)
            h5_file.attrs['created_time'] = time.strftime('%Y-%m-%d %H:%M:%S')

        print(f"已保存所有股票代码到 stock_ticker.h5 文件")
        print(f"上海交易所: {len(sh_tickers)} 只股票")
        print(f"深圳交易所: {len(sz_tickers)} 只股票")
        print(f"总计: {len(sh_tickers) + len(sz_tickers)} 只股票")

    except Exception as e:
        print(f"保存文件时出错: {e}")

def main():
    """
    主函数：获取上海和深圳的股票代码并保存到h5文件
    """
    # saveTickers()
    # stock_zh_a_hist_df = ak.stock_zh_a_hist(symbol="000001", period="daily", start_date="20170301", end_date='20240528',adjust="qfq")
    # print(stock_zh_a_hist_df)

    import yfinance as yf

    # 下载浦发银行（600000.SS）的历史数据
    DOW_30_TICKER = [
        "AXP",
        # "AMGN",
        # "AAPL",
        # "BA",
        # "CAT",
        # "CSCO",
        # "CVX",
        # "GS",
        # "HD",
        # "HON",
        # "IBM",
        # "INTC",
        # "JNJ",
        # "KO",
        # "JPM",
        # "MCD",
        # "MMM",
        # "MRK",
        # "MSFT",
        # "NKE",
        # "PG",
        # "TRV",
        # "UNH",
        # "CRM",
        # "VZ",
        # "V",
        # "WBA",
        # "WMT",
        # "DIS",
        # "DOW",
    ]
    start_date = "2023-11-01"
    end_date = "2023-12-31"
    # for tic in DOW_30_TICKER:
    #     try:
    #         # 前复权 下载数据
    #         data = yf.download(tic, start=start_date, end=end_date)
    #         print(f"已下载 {tic} 的历史数据")
    #         print(data.head())
    #     except YFTzMissingError:
    #         print(f"{tic} 的历史数据缺少时区信息，请检查数据源")
    #         continue


    start =  pd.Timestamp(start_date)
    formatted_date = start.strftime("%Y-%m-%d")

    print(formatted_date)  # 输出: 2020-01-06
    # 打印start格式 yyyy-mm-dd 的字符串




    start_date = datetime(2020, 1, 6)
    end_date = datetime(2024, 1, 31)
    delta = (end_date - start_date).days
    print(delta)

    a=[1,2]+np.zeros(3).tolist()
    print(a)



if __name__ == "__main__":
    main()